#include "yolo_complete.h"

namespace sp::Model::Yolo{

void registe_yolo_complete(){
    yolo_complete::registe();
}

yolo_complete::yolo_complete(std::vector<std::string> v): PythonFunction("/home/lx/SmartPipe/src/core/functions/Model/Yolo/yolo/", "/home/lx/SmartPipe/src/core/functions/Model/Yolo/yolo/repo/yolov5/", "yolo_complete", "yolo", v), GpuFunction(){
    assert(v.size() == 0);
    name = "yolo_complete";
    class_name = "Model-Yolo-yolo_complete";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT}, RECT));
}

yolo_complete::yolo_complete(): PythonFunction(), GpuFunction(){
    name = "yolo_complete";
    class_name = "Model-Yolo-yolo_complete";
    pModulePath = "/home/lx/SmartPipe/src/core/functions/Model/Yolo/yolo/";
    pModuleHomePath = "/home/lx/SmartPipe/src/core/functions/Model/Yolo/yolo/repo/yolov5/";
    pModuleName = "yolo_complete";
    pClassName = "yolo";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({MAT}, RECT));
}

yolo_complete::~yolo_complete(){

}

void yolo_complete::registe(){
    Function::Register("Model-Yolo-yolo_complete", createObject);
    Function::Register("Model-Yolo-yolo_complete", createObject2);
}

Function* yolo_complete::createObject(std::vector<std::string> params){
    return new yolo_complete(params);
}

Function* yolo_complete::createObject2(){
    return new yolo_complete();
}

void yolo_complete::start(){
    // 构建实例
    PythonFunction::defaultStart();
    pInstance = PyObject_CallObject(pConstruct, nullptr);
    assert(pInstance != nullptr);
    PyObject_CallMethod(pInstance, "start", "");
}

bool yolo_complete::waitForResource(){
    return true;
}

/*
    vector<vector<Mat*>> -> [Mat,Mat,...] 
        yolo_inference: [Mat,Mat,...] -> [[Rect,Rect,...],[Rect,Rect,...],...]
    [[Rect,Rect,...],[Rect,Rect,...],...] -> vector<Head*> (Head* -> Rect* -> Rect* -> ...)
*/
void yolo_complete::process(std::vector<std::vector<Data*>>& data_input, std::vector<Data*>& data_output){
    assert(data_input.size() != 0);
    assert(data_input[0].size() == 1);
    // 将data_input中的数据提取出来并转成numpy矩阵
    PyObject* data = convertToPython(data_input);
    // 调用python函数
    pArgs = PyTuple_New(1);
    PyTuple_SetItem(pArgs, 0, data);
    pReturn = PyObject_CallMethod(pInstance, "inference", "O", pArgs);
    // 将结果解析成Data并赋值给data_output
    convertToCpp(pReturn, data_input, data_output, 1); // 输出在第一层链起来
    // 释放内存
    for(int i = 0; i < data_input.size(); i++){
        assert(data_input[i][0]->type == MAT);
        if(data_input[i][0]->subRefCount() == 0){
            executor->freeAll(data_input[i][0]);
        }
    }
}

bool yolo_complete::releaseResource(){
    return true;
}

void yolo_complete::finish(){
    PyObject_CallMethod(pInstance, "finish", "");
    PythonFunction::defaultFinish();
}

void yolo_complete::copy(Function* other){
    Function::defaultCopy(other);
}

};